package com.neu.system.controller;

import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.client.RestTemplate;
import java.util.*;
import org.springframework.web.multipart.MultipartFile;
import java.io.IOException;
import java.net.HttpURLConnection;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.util.Base64;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import java.io.BufferedReader;
import java.io.InputStreamReader;

@RestController
@RequestMapping("/api/image")
public class ImageRecognitionController {

    private static final String API_KEY = "f9738f83ab1f4480a8f1faeaba2fbca3.xEgcfpiGDnaJikgk";
    private static final String API_URL = "https://open.bigmodel.cn/api/paas/v4/chat/completions";

    @PostMapping("/recognize")
    public ResponseEntity<?> recognizeImage(@RequestBody Map<String, String> body) {
        String base64 = body.get("image_base64");
        String question = body.getOrDefault("question", "请描述这张图片");

        // 官方推荐body结构
        List<Object> contentList = new ArrayList<>();
        Map<String, Object> imageObj = new HashMap<>();
        imageObj.put("type", "image_url");
        Map<String, Object> imageUrl = new HashMap<>();
        imageUrl.put("url", "data:image/jpeg;base64," + base64);
        imageObj.put("image_url", imageUrl);
        contentList.add(imageObj);
        Map<String, Object> textObj = new HashMap<>();
        textObj.put("type", "text");
        textObj.put("text", question);
        contentList.add(textObj);
        Map<String, Object> userMsg = new HashMap<>();
        userMsg.put("role", "user");
        userMsg.put("content", contentList);
        List<Map<String, Object>> messages = new ArrayList<>();
        messages.add(userMsg);
        Map<String, Object> payload = new HashMap<>();
        payload.put("model", "glm-4v-plus-0111");
        payload.put("messages", messages);
        // 发送请求
        RestTemplate restTemplate = new RestTemplate();
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.set("Authorization", "Bearer " + API_KEY);
        HttpEntity<Map<String, Object>> entity = new HttpEntity<>(payload, headers);
        try {
            ResponseEntity<String> response = restTemplate.postForEntity(API_URL, entity, String.class);
            return ResponseEntity.status(response.getStatusCode()).body(response.getBody());
        } catch (Exception e) {
            return ResponseEntity.status(500).body(Collections.singletonMap("error", e.getMessage()));
        }
    }

    @PostMapping(value = "/recognize/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter recognizeImageStream(
            @RequestParam("file") MultipartFile file,
            @RequestParam(value = "question", required = false, defaultValue = "请描述这张图片") String question
    ) throws IOException {
        SseEmitter emitter = new SseEmitter(300_000L);
        new Thread(() -> {
            try {
                // 1. 图片转base64
                String base64Image = Base64.getEncoder().encodeToString(file.getBytes());
                // 官方推荐body结构
                List<Object> contentList = new ArrayList<>();
                Map<String, Object> imageObj = new HashMap<>();
                imageObj.put("type", "image_url");
                Map<String, Object> imageUrl = new HashMap<>();
                imageUrl.put("url", "data:image/jpeg;base64," + base64Image);
                imageObj.put("image_url", imageUrl);
                contentList.add(imageObj);
                Map<String, Object> textObj = new HashMap<>();
                textObj.put("type", "text");
                textObj.put("text", question);
                contentList.add(textObj);
                Map<String, Object> userMsg = new HashMap<>();
                userMsg.put("role", "user");
                userMsg.put("content", contentList);
                List<Map<String, Object>> messages = new ArrayList<>();
                messages.add(userMsg);
                Map<String, Object> payload = new HashMap<>();
                payload.put("model", "glm-4v-plus-0111");
                payload.put("messages", messages);
                payload.put("stream", true);
                String json = new com.fasterxml.jackson.databind.ObjectMapper().writeValueAsString(payload);
                // 3. 发送流式请求
                URL url = new URL(API_URL);
                HttpURLConnection conn = (HttpURLConnection) url.openConnection();
                conn.setRequestMethod("POST");
                conn.setRequestProperty("Authorization", "Bearer " + API_KEY);
                conn.setRequestProperty("Content-Type", "application/json");
                conn.setDoOutput(true);
                conn.getOutputStream().write(json.getBytes(StandardCharsets.UTF_8));
                BufferedReader reader = new BufferedReader(new InputStreamReader(conn.getInputStream(), StandardCharsets.UTF_8));
                String line;
                while ((line = reader.readLine()) != null) {
                    if (line.startsWith("data:")) {
                        String data = line.substring(5).trim();
                        if ("[DONE]".equals(data)) break;
                        emitter.send(data);
                    }
                }
                emitter.complete();
            } catch (Exception e) {
                try { emitter.send("[ERROR] " + e.getMessage()); } catch (Exception ignore) {}
                emitter.completeWithError(e);
            }
        }).start();
        return emitter;
    }
} 